from langchain.chat_models import init_chat_model
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser  # 新版本解析器
from pydantic import BaseModel, Field  # 用Pydantic定义结构（替代ResponseSchema）
from langchain_core.runnables import RunnableLambda




import os
from dotenv import load_dotenv
load_dotenv()

def debug_print(x):
    print('中间结果（新闻正文）：', x)
    return x

debug_node = RunnableLambda(debug_print) # 创建一个打印中间结果的节点，用来打印中间结果

# 第一步：根据标题生成新闻正文
news_gen_prompt = PromptTemplate.from_template(
    "请根据以下新闻标题撰写一段简短的新闻内容（100字以内）：\n\n标题：{title}"
)

model = init_chat_model(
    model="Qwen/Qwen3-8B",
    model_provider="openai",
    base_url="https://api.siliconflow.cn/v1/",
    api_key=os.environ.get('OPENAI_API_KEY'),
)

# 第一个子链：生成新闻内容
news_chain = news_gen_prompt | model

# 第二步：用Pydantic模型定义需要提取的结构化字段（替代ResponseSchema）
class NewsInfo(BaseModel):
    time: str = Field(description="事件发生时间")
    location: str = Field(description="事件发生地点")
    event: str = Field(description="发生的具体事件")

# 初始化JsonOutputParser，作用是将模型输出的文本转换为JSON
parser = JsonOutputParser(pydantic_object=NewsInfo)

# 定义提取信息的提示词模板（格式说明通过parser生成）
summary_prompt = PromptTemplate(
    template="请将新闻内容进行总结，并返回一个JSON格式的字符串，格式为：\n{format_instructions}\n\n新闻内容：{news_content}",
    input_variables=["news_content"],
    partial_variables={"format_instructions": parser.get_format_instructions()}
)

# 第二个子链：提取结构化信息（用JsonOutputParser解析）
summary_chain = summary_prompt | model | parser

# 组合成复合链（先生成新闻，再提取信息）
full_chain = news_chain | debug_node  | summary_chain

# 调用复合链
result = full_chain.invoke({"title": "北京大学发文悼念杨振宁"})
print(result)
